package edu.washington.ling.syntric.proc;

import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import org.apache.commons.configuration.ConfigurationException;
import org.apache.commons.configuration.XMLConfiguration;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.config.Configuration;
import edu.washington.ling.syntric.io.*;
import edu.washington.ling.syntric.struct.error.*;
import edu.washington.ling.syntric.struct.rule.*;
import edu.washington.ling.syntric.struct.tree.Node;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.ConvertUtils;
import edu.washington.ling.syntric.util.MergeUtils;
import edu.washington.ling.syntric.util.Utils;

public class ExtractBinaryRules {

	private static ExtractBinaryRulesOptions opts = new ExtractBinaryRulesOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	
	
	private static List<Node> getMaxChain(Node root, List<Node> chain, Map<Node, Set<Node>> rootHeadMap) throws Exception {
		if (new HashSet<Node>(chain).contains(root)) {
			throw new Exception(String.format("Warning: cycle found. Rule root: %s, Existing chain: %s.", root.getLabel(), formatNodes(chain)));
		}
		List<Node> maxChain = new LinkedList<Node>(chain);
		maxChain.add(root);
		if (rootHeadMap.containsKey(root)) {
			List<Node> maxHeadChain = new LinkedList<Node>();
			List<Node> headChain;
			for (Node headChild: rootHeadMap.get(root)) {
				headChain = getMaxChain(headChild, maxChain, rootHeadMap);
				if (headChain.size() > maxHeadChain.size()) {
					maxHeadChain = headChain;
				}
			}
			maxChain = maxHeadChain;
		}
		return maxChain;
	}
	private static Map<Token, List<BiRule>> getHeadRuleMap(List<BiRule> rules) {
		Map<Token, List<BiRule>> headRuleMap = new HashMap<Token, List<BiRule>>();
		for (BiRule rule: rules) {
			BiDSPattern dsPatt = rule.getDSPattern();
			Token head = dsPatt.getDSHead();
			if (headRuleMap.containsKey(head)) {
				headRuleMap.get(head).add(rule);
			}
			else {
				List<BiRule> headRules = new LinkedList<BiRule>();
				headRules.add(rule);
				headRuleMap.put(head, headRules);
			}
		}
		return headRuleMap;
	}
	private static Map<Node, Set<Node>> getRootHeadMap(List<BiRule> rules) {
		
		Map<Node, Set<Node>> rootHeadMap = new HashMap<Node, Set<Node>>();
		for (BiRule rule: rules) {
			BiPSPattern psPatt = rule.getPSPattern();
			Node root = psPatt.getRoot();
			Node headChild = psPatt.getHeadChild();
			if (headChild.equals(root)) {
				continue;
			}
			if (rootHeadMap.containsKey(root)) {
				rootHeadMap.get(root).add(headChild);
			}
			else {
				Set<Node> headChildren = new HashSet<Node>();
				headChildren.add(headChild);
				rootHeadMap.put(root, headChildren);
			}
		}
		return rootHeadMap;
	}
	private static Set<Node> getChainRoots(List<BiRule> rules) {
		Set<Node> roots = new HashSet<Node>();
		Set<Node> headChildren = new HashSet<Node>();
		for (BiRule rule: rules) {
			BiPSPattern psPatt = rule.getPSPattern();
			Node root = psPatt.getRoot();
			roots.add(root);
			Node headChild = psPatt.getHeadChild();
			if (!headChild.equals(root)) {
				headChildren.add(headChild);
			}
		}
		roots.removeAll(headChildren);
		return roots;
	}
	private static void checkStackingAmbiguity(Map<Token, List<BiRule>> headRuleMap) {
		int adjunctCount = 0;
		int adjunctsWithStackingAmbiguity = 0;
		Set<Node> roots;
		Map<Node, Set<Node>> unorderedDependents;
		Map<Node, Set<Node>> rightDependents;
		Map<Node, Set<Node>> leftDependents;
		for (Token head: headRuleMap.keySet()) {
			List<BiRule> rules = headRuleMap.get(head);
			roots = new HashSet<Node>();
			unorderedDependents = new HashMap<Node, Set<Node>>();
			rightDependents = new HashMap<Node, Set<Node>>();
			leftDependents = new HashMap<Node, Set<Node>>();
			for (BiRule rule: rules) {
				BiDSPattern dsPatt = rule.getDSPattern();
				BiPSPattern psPatt = rule.getPSPattern();
				Node root = psPatt.getRoot();
				Node headChild = psPatt.getHeadChild();
				Node depChild = psPatt.getDepChild();
				if (!headChild.equals(root)) {
					continue;
				}
				adjunctCount++;
				roots.add(root);
				if (dsPatt.hasHeadOrder()) {
					if(dsPatt.isHeadFinal()) {
						if (leftDependents.containsKey(root)) {
							leftDependents.get(root).add(depChild);
						}
						else {
							leftDependents.put(root, new HashSet<Node>());
							leftDependents.get(root).add(depChild);
						}
					}
					else {
						if (rightDependents.containsKey(root)) {
							rightDependents.get(root).add(depChild);
						}
						else {
							rightDependents.put(root, new HashSet<Node>());
							rightDependents.get(root).add(depChild);
						}
					}
				}
				else {
					if (unorderedDependents.containsKey(root)) {
						unorderedDependents.get(root).add(depChild);
					}
					else {
						unorderedDependents.put(root, new HashSet<Node>());
						unorderedDependents.get(root).add(depChild);
					}
				}
			}
			for (Node root: roots) {
				if (leftDependents.containsKey(root) && rightDependents.containsKey(root)) {
					String errStr1 = String.format("%s over head %s has multiple stacking possibilities.", 
							                      root.getLabel(), head.getTag());
					String errStr2 = String.format("Left dependents: %s",formatNodes(leftDependents.get(root)));
					String errStr3 = String.format("Right dependents: %s",formatNodes(rightDependents.get(root)));
					System.err.println(errStr1);
					System.err.println(errStr2);
					System.err.println(errStr3);
					adjunctsWithStackingAmbiguity++;
				}
				if (unorderedDependents.containsKey(root)) {
					String errStr1 = String.format("%s over head %s has multiple stacking possibilities.", 
													root.getLabel(), head.getTag());
					String errStr2 = String.format("Unordered dependents: %s",formatNodes(unorderedDependents.get(root)));
					System.err.println(errStr1);
					System.err.println(errStr2);
					adjunctsWithStackingAmbiguity++;
				}
			}
		}
		System.out.println("Stacking ambiguity:");
		System.out.println(String.format("%d/%d adjuncts have potential stacking ambiguity.", 
				           adjunctsWithStackingAmbiguity, adjunctCount));
		
	}
	private static void checkProjectionOrders(Map<Token, List<BiRule>> headRuleMap) {
		int orderCount = 0;
		int ordersWithCycles = 0;
		for (Token head: headRuleMap.keySet()) {
			List<BiRule> rules = headRuleMap.get(head);
			Map<Node, Set<Node>> rootHeadMap = getRootHeadMap(rules);
			Set<Node> chainRoots = getChainRoots(rules);
			System.out.println(String.format("\nProjection orders for head=%s:", head.getTag()));
			for(Node chainRoot: chainRoots) {
				orderCount++;
				try {
					List<Node> maxChain = getMaxChain(chainRoot, new LinkedList<Node>(), rootHeadMap);
					System.out.println(formatNodes(maxChain));
				}
				catch(Exception e) {
					System.err.println(e.getMessage());
					ordersWithCycles++;
				}
			}
		}
		System.out.println("Projection order check complete.");
		System.out.println(String.format("%d/%d projection orders have cycles.", ordersWithCycles, orderCount));
	}
	private static void checkProjections(List<BiRule> rules) {
		Map<Token, List<BiRule>> headRuleMap = getHeadRuleMap(rules);
		checkStackingAmbiguity(headRuleMap);
		//checkProjectionOrders(headRuleMap); **Should be output to file when this is a separate step**
	}
	private static String formatNodes(Collection<Node> nodes) {
		LinkedList<String> labelList = new LinkedList<String>();
		for (Node node: nodes) {
			labelList.add(node.getLabel());
		}
		return Utils.join(labelList, " ");
	}
	private static List<BiRule> getAmbiguousRules (List<BiRule> rules) {
		List<BiRule> ambigRules = new LinkedList<BiRule>();
		boolean ruleIsUnambig;
		for (int i = 0; i < rules.size(); i++) {
			BiRule rule = rules.get(i);
			ruleIsUnambig = true;
			if (i < rules.size() - 1) {
				BiRule nextRule = rules.get(i+1);
				if (rule.getDSPattern().equals(nextRule.getDSPattern()))
					ruleIsUnambig = false;
			}
			if (i > 0) {
				BiRule prevRule = rules.get(i-1);
				if (rule.getDSPattern().equals(prevRule.getDSPattern()))
					ruleIsUnambig = false;
			}
			if (!ruleIsUnambig) {
				ambigRules.add(rules.get(i));
			}
		}
		return ambigRules;
	}
	private static void printRuleStatistics(Map<BiDSPattern, List<BiPSPattern>> ruleMap) {
		int numAmbig = 0;
		int numBiAmbig = 0;
		int numTriAmbig = 0;
		int numMoreAmbig = 0;
		for (BiDSPattern dsPatt: ruleMap.keySet()) {
			Integer count = ruleMap.get(dsPatt).size();
			if (count > 1) {
				numAmbig += 1;
			}
			if (count == 2) {
				numBiAmbig += 1;
			}
			else if (count == 3) {
				numTriAmbig += 1;
			}
			else if (count > 3) {
				numMoreAmbig += 1;
			}
		}
		System.out.println("Ambiguity Statistics:");
		System.out.println(String.format("A total of %d DS patterns were extracted.", ruleMap.size()));
		System.out.println(String.format("A total of %d were ambiguous.", numAmbig));
		System.out.println(String.format("A total of %d were two-ways ambiguous.", numBiAmbig));
		System.out.println(String.format("A total of %d were three-ways ambiguous.", numTriAmbig));
		System.out.println(String.format("A total of %d were more ambiguous.", numMoreAmbig));
	}
	private static int getTotalExampleCount(BiDSPattern dsPatt, List<BiPSPattern> psPatts, 
			Map<BiRule, List<Example>> ruleExMap) {
		int totalCount = 0;
		for (BiPSPattern psPatt: psPatts) {
			totalCount += ruleExMap.get(new BiRule(dsPatt, psPatt)).size();
		}
		return totalCount;
	}
	public static void printToFile(Map<BiDSPattern, List<BiPSPattern>> ruleMap, 
			Map<BiRule, List<Example>> ruleExMap, String path, String format) {
		printToFile(ruleMap, ruleExMap, path, format, false);
	}
	public static void printToFile(Map<BiDSPattern, List<BiPSPattern>> ruleMap, 
			Map<BiRule, List<Example>> ruleExMap, String path, String format, boolean ambigOnly) {
		try {
			RuleWriter rw = RuleWriter.newWriter(format);
			RuleFileWriter rfw = RuleFileWriter.newWriter(path, format);
			List<BiDSPattern> dsPatts = new LinkedList<BiDSPattern>(ruleMap.keySet());
			Collections.sort(dsPatts);
			if (rfw instanceof RuleAndExampleFileWriter) {
				RuleAndExampleFileWriter refw = (RuleAndExampleFileWriter) rfw;
				String[] includeLayers = new String[] {"PS", "DS-Const"};
				for (BiDSPattern dsPatt: dsPatts) {
					List<BiPSPattern> psPatts = ruleMap.get(dsPatt);
					if (psPatts.size() > 1) {
						int totalRules = psPatts.size();
						int totalExamples = getTotalExampleCount(dsPatt, psPatts, ruleExMap);
						refw.writeHeaderForAmbigRule(rw.writeBinaryDSPattern(dsPatt), totalRules);
						for (BiPSPattern psPatt: psPatts) {
							BiRule rule = new BiRule(dsPatt, psPatt);
							List<Example> exrs = ruleExMap.get(rule);
							int writeLimit = 1;
							if ((float) exrs.size() / totalExamples < 0.33) writeLimit = exrs.size();
							refw.writeRuleAndExampleReps(rw.writeBinaryRule(rule), exrs, includeLayers, writeLimit);
						}
						refw.writeFooterForAmbigRule();
					} else if (psPatts.size() == 1 && ambigOnly == false) {
						BiRule rule = new BiRule(dsPatt, psPatts.get(0));
						List<Example> exrs = ruleExMap.get(rule);
						refw.writeRuleAndExampleReps(rw.writeBinaryRule(rule), exrs, includeLayers, 1);
					}
				}
			}
			else {
				for (BiDSPattern dsPatt: dsPatts) {
					for (BiPSPattern psPatt: ruleMap.get(dsPatt)) {
						BiRule rule = new BiRule(dsPatt, psPatt);
						List<Example> exrs = new LinkedList<Example>(ruleExMap.get(rule));
						rfw.writeRuleRep(rw.writeBinaryRule(rule), exrs.size());
					}
				}
			}
			rfw.close();
		} catch(Exception e) {
			pHandler.handle(e, "Printing of binary rules failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}

	public static BiRule extractBinaryRule(Token dsDependent, Token dsHead, Tree tree, Configuration config) {

		BiDSPattern dsPatt = ConvertUtils.extractBinaryDSPattern(dsDependent, dsHead, 
										config.getDepFeats(), config.getHeadFeats(),
										config.useExtPOS());
		BiPSPattern psPatt = ConvertUtils.extractBinaryPSPattern(dsDependent, dsHead);
		if(config.extractOrder()) {
			boolean headFinal = true;
			if (tree.getTokenPosition(dsDependent) > tree.getTokenPosition(dsHead)) {
				headFinal = false;
			}
			dsPatt.setHeadFinal(headFinal);
		}
		return new BiRule(dsPatt, psPatt);
	}
	public static Set<BiRule> extractBinaryRules(Tree tree, Configuration config) {
		Set<BiRule> biRules = new HashSet<BiRule>();
		for(Token t: tree.getTokenList()) {
			if (t.hasDSHead()) {
				Token head = t.getDSHead();
				biRules.add(extractBinaryRule(t, head, tree, config));
			}
		}
		return biRules;
	}
	private static Map<BiRule, List<Example>> extractRulesFromInputFile(Configuration config) {
		Map<BiRule, List<Example>> ruleExMap = new HashMap<BiRule, List<Example>>();
		int numExamples = 0;
		int numCompat = 0;
		try {
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			TreeReader dsr = TreeReader.newReader(opts.dsTreeFormat);
			TreeReader psr = TreeReader.newReader(opts.psTreeFormat);
			Example exr = null;
			while ((exr = exf.getNextExample()) != null) {
				numExamples++;
				if (ConvertUtils.isCompatible(exr)) {
					numCompat++;
					try {
						Tree ds = dsr.readTree(exr.getTreeRep("DS-Const"));
						Tree ps = psr.readTree(exr.getTreeRep("PS"));
						Tree tree = MergeUtils.mergeDSandPSTrees(ds, ps);
						Set<BiRule> biRules = extractBinaryRules(tree, config);
						for (BiRule biRule: biRules) {
							if (!ruleExMap.containsKey(biRule)) {
								ruleExMap.put(biRule, new LinkedList<Example>());
							}
							ruleExMap.get(biRule).add(exr);
						}
						
					}
					catch (TreeMergeException e) {
						String errStr = String.format("Could not merge representations for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
					catch (UnreadableRepresentationException e) {
						String errStr = String.format("Could not process representation for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
				}
				
			}
			System.out.println("Extraction of binary rules was successful.");
			System.out.println(String.format("A total of %d rules were extracted from %d compatible tree pairs.", 
											 ruleExMap.size(), numCompat));
			System.out.println(String.format("Processed %d total examples", numExamples));
		} catch(Exception e) {
			pHandler.handle(e, "Extraction of binary rules failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
		return ruleExMap;
	}
	

	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			Configuration config;
			if (opts.configFilePath != null) {
				XMLConfiguration xmlConfig = new XMLConfiguration(opts.configFilePath);
				config = Configuration.init(xmlConfig);
			} else {
				config = Configuration.initDefault();
			}
			Map<BiRule, List<Example>> ruleExMap = extractRulesFromInputFile(config);
			Map<BiDSPattern, List<BiPSPattern>> ruleMap = new HashMap<BiDSPattern, List<BiPSPattern>>();
			for (BiRule rule: ruleExMap.keySet()) {
				if (!ruleMap.containsKey(rule.getDSPattern())) {
					ruleMap.put(rule.getDSPattern(), new LinkedList<BiPSPattern>());
				}
				ruleMap.get(rule.getDSPattern()).add(rule.getPSPattern());
			}
			printRuleStatistics(ruleMap);
			checkProjections(new LinkedList<BiRule>(ruleExMap.keySet()));
			printToFile(ruleMap, ruleExMap, opts.outputFilePath, "text");
			if (opts.latexOutputFilePath != null) {
				printToFile(ruleMap, ruleExMap, opts.latexOutputFilePath, "latex", true);
			}
			
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("Usage: ExtractBinaryRules");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(1);
		}
		catch(ConfigurationException e) {
			String errStr = String.format("configuration file \'%s\' could not be read.", opts.configFilePath);
			pHandler.handle(e, errStr);
			System.exit(1);
		}
	}
}
